"""
均线金叉死叉算法验证工具 (简化版)
功能：验证趋势追踪止损中使用的均线上穿下穿算法
输出：生成一个列表，标记每个K线周期的金叉(上穿)、死叉(下穿)或无信号状态
可视化：使用ECharts生成交互式K线图和双均线副图，在副图上用上下箭头标注金死叉
特点：
1. 主图显示K线走势
2. 副图显示双均线(Trading Line和Short Line)
3. 金叉(上穿)用绿色上箭头标注
4. 死叉(下穿)用红色下箭头标注
5. 鼠标悬停可查看详细信息
"""

# 全局配置参数 - 用户可根据需要调整
MAX_KLINES_TO_DISPLAY = 300  # 最大显示K线数量，默认300根
DEFAULT_YEARS = 2.0  # 默认获取最近几年的数据

import numpy as np
import pandas as pd
import akshare as ak  # 正确导入 akshare 库
from datetime import datetime
import argparse
from dateutil.relativedelta import relativedelta
from MyTT import REF, EMA

# 新增博易大师技术指标算法
def calculate_boeyi_indicators(klines):
    """
    计算博易大师技术指标，包含Q值、操盘线和空头线
    """
    close = klines.close
    open_price = klines.open
    high = klines.high
    low = klines.low
    
    # 计算Q值
    Q = (3*close + low + open_price + high) / 6
    Q = pd.Series(Q)
    
    # 手动实现 REF 函数
    def ref(series, n=1):
        return series.shift(n)
    
    # 计算操盘线（26期Q值加权平均）
    terms = [
        26 * Q,
        25 * ref(Q, 1),
        24 * ref(Q, 2),
        23 * ref(Q, 3),
        22 * ref(Q, 4),
        21 * ref(Q, 5),
        20 * ref(Q, 6),
        19 * ref(Q, 7),
        18 * ref(Q, 8),
        17 * ref(Q, 9),
        16 * ref(Q, 10),
        15 * ref(Q, 11),
        14 * ref(Q, 12),
        13 * ref(Q, 13),
        12 * ref(Q, 14),
        11 * ref(Q, 15),
        10 * ref(Q, 16),
        9 * ref(Q, 17),
        8 * ref(Q, 18),
        7 * ref(Q, 19),
        6 * ref(Q, 20),
        5 * ref(Q, 21),
        4 * ref(Q, 22),
        3 * ref(Q, 23),
        2 * ref(Q, 24),
        1 * ref(Q, 25)
    ]
    
    # 计算操盘线
    trading_line_numerator = sum(terms)
    trading_line_denominator = sum(range(1, 27))  # 1+2+...+26
    trading_line = trading_line_numerator / trading_line_denominator
    
    # 计算空头线（52期Q值加权平均）
    terms = [
        52 * Q,
        51 * ref(Q, 1),
        50 * ref(Q, 2),
        49 * ref(Q, 3),
        48 * ref(Q, 4),
        47 * ref(Q, 5),
        46 * ref(Q, 6),
        45 * ref(Q, 7),
        44 * ref(Q, 8),
        43 * ref(Q, 9),
        42 * ref(Q, 10),
        41 * ref(Q, 11),
        40 * ref(Q, 12),
        39 * ref(Q, 13),
        38 * ref(Q, 14),
        37 * ref(Q, 15),
        36 * ref(Q, 16),
        35 * ref(Q, 17),
        34 * ref(Q, 18),
        33 * ref(Q, 19),
        32 * ref(Q, 20),
        31 * ref(Q, 21),
        30 * ref(Q, 22),
        29 * ref(Q, 23),
        28 * ref(Q, 24),
        27 * ref(Q, 25),
        26 * ref(Q, 26),
        25 * ref(Q, 27),
        24 * ref(Q, 28),
        23 * ref(Q, 29),
        22 * ref(Q, 30),
        21 * ref(Q, 31),
        20 * ref(Q, 32),
        19 * ref(Q, 33),
        18 * ref(Q, 34),
        17 * ref(Q, 35),
        16 * ref(Q, 36),
        15 * ref(Q, 37),
        14 * ref(Q, 38),
        13 * ref(Q, 39),
        12 * ref(Q, 40),
        11 * ref(Q, 41),
        10 * ref(Q, 42),
        9 * ref(Q, 43),
        8 * ref(Q, 44),
        7 * ref(Q, 45),
        6 * ref(Q, 46),
        5 * ref(Q, 47),
        4 * ref(Q, 48),
        3 * ref(Q, 49),
        2 * ref(Q, 50),
        1 * ref(Q, 51)
    ]
    
    # 计算空头线
    short_line_numerator = sum(terms)
    short_line_denominator = sum(range(1, 53))  # 1+2+...+52
    short_line = short_line_numerator / short_line_denominator
    
    return trading_line, short_line

# 计算交易信号
def calculate_trading_signals(klines):
    """
    计算交易信号：
    - 操盘线上穿空头线：买入信号（金叉）
    - 操盘线下穿空头线：卖出信号（死叉）
    """
    # 计算操盘线和空头线
    trading_line, short_line = calculate_boeyi_indicators(klines)
    
    # 创建结果DataFrame
    result = pd.DataFrame({
        'datetime': klines.datetime,
        'open': klines.open,
        'high': klines.high,
        'low': klines.low,
        'close': klines.close,
        'trading_line': trading_line,
        'short_line': short_line
    })
    
    # 计算信号
    # 上穿：当前交易线在空头线之上，且前一周期在空头线之下
    # 下穿：当前交易线在空头线之下，且前一周期在空头线之上
    result['signal'] = None
    
    for i in range(1, len(result)):
        # 确保两个周期的值都不是NaN
        if (pd.notna(result['trading_line'].iloc[i-1]) and 
            pd.notna(result['short_line'].iloc[i-1]) and
            pd.notna(result['trading_line'].iloc[i]) and 
            pd.notna(result['short_line'].iloc[i])):
            
            # 上穿信号（金叉）
            if (result['trading_line'].iloc[i-1] <= result['short_line'].iloc[i-1] and 
                result['trading_line'].iloc[i] > result['short_line'].iloc[i]):
                result.loc[result.index[i], 'signal'] = '上穿'
            
            # 下穿信号（死叉）
            elif (result['trading_line'].iloc[i-1] >= result['short_line'].iloc[i-1] and 
                  result['trading_line'].iloc[i] < result['short_line'].iloc[i]):
                result.loc[result.index[i], 'signal'] = '下穿'
    
    # 打印操盘线的前几个值，用于调试
    print("操盘线值：", trading_line.head(3))
    
    return result

def generate_html(result, symbol, period_name):
    """
    生成HTML文件，使用原生ECharts，包含三图布局（一个主图和两个副图）
    """
    # 准备数据
    date_list = result['datetime'].astype(str).tolist()
    k_data = result[['open', 'close', 'low', 'high']].values.tolist()
    trading_line_data = result['trading_line'].tolist()
    short_line_data = result['short_line'].tolist()
    
    # 准备金叉和死叉数据
    up_cross = result[result['signal'] == '上穿']
    down_cross = result[result['signal'] == '下穿']
    
    up_cross_x = up_cross.index.tolist()
    up_cross_y = up_cross['trading_line'].tolist()
    up_cross_dates = up_cross['datetime'].astype(str).tolist()
    
    down_cross_x = down_cross.index.tolist()
    down_cross_y = down_cross['trading_line'].tolist()
    down_cross_dates = down_cross['datetime'].astype(str).tolist()
    
    # 计算操盘线和空头线的差值
    diff_data = []
    for i in range(len(result)):
        if pd.notna(result['trading_line'].iloc[i]) and pd.notna(result['short_line'].iloc[i]):
            diff = result['trading_line'].iloc[i] - result['short_line'].iloc[i]
            diff_data.append(diff)
        else:
            diff_data.append(None)
    
    # 创建HTML内容
    html_content = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>{symbol} {period_name} 金叉死叉分析</title>
    <script src="https://cdn.jsdelivr.net/npm/echarts@5/dist/echarts.min.js"></script>
    <style>
        body {{
            margin: 0;
            padding: 20px;
            font-family: Arial, sans-serif;
        }}
        .container {{
            width: 100%;
            max-width: 1200px;
            margin: 0 auto;
        }}
        .chart-container {{
            width: 100%;
            height: 1000px;
            margin-bottom: 20px;
            border: 1px solid #ddd;
            box-shadow: 0 0 10px rgba(0,0,0,0.1);
        }}
        h1 {{
            text-align: center;
            margin-bottom: 20px;
        }}
        .info {{
            margin: 20px 0;
            padding: 15px;
            background-color: #f8f9fa;
            border-radius: 5px;
            border-left: 5px solid #007bff;
        }}
        .stats {{
            display: flex;
            justify-content: space-around;
            margin: 20px 0;
            padding: 10px;
            background-color: #f8f9fa;
            border-radius: 5px;
        }}
        .stat-item {{
            text-align: center;
        }}
        .stat-value {{
            font-size: 24px;
            font-weight: bold;
        }}
        .gold-cross {{
            color: #28a745;
        }}
        .death-cross {{
            color: #dc3545;
        }}
    </style>
</head>
<body>
    <div class="container">
        <h1>{symbol} {period_name} 金叉死叉分析</h1>
        
        <div class="info">
            <p><strong>金叉信号</strong>：操盘线上穿空头线，形成买入信号，用绿色三角形标注</p>
            <p><strong>死叉信号</strong>：操盘线下穿空头线，形成卖出信号，用红色三角形标注</p>
            <p><strong>差值图</strong>：显示操盘线与空头线的差值，正值区域为红色(多头)，负值区域为绿色(空头)</p>
        </div>
        
        <div class="stats">
            <div class="stat-item">
                <div class="stat-label">金叉信号数量</div>
                <div class="stat-value gold-cross">{len(up_cross)}</div>
            </div>
            <div class="stat-item">
                <div class="stat-label">死叉信号数量</div>
                <div class="stat-value death-cross">{len(down_cross)}</div>
            </div>
            <div class="stat-item">
                <div class="stat-label">总信号数量</div>
                <div class="stat-value">{len(up_cross) + len(down_cross)}</div>
            </div>
            <div class="stat-item">
                <div class="stat-label">信号占比</div>
                <div class="stat-value">{((len(up_cross) + len(down_cross)) / len(result) * 100):.2f}%</div>
            </div>
        </div>
        
        <div id="main-chart" class="chart-container"></div>
    </div>

    <script>
        // 初始化图表
        var chartDom = document.getElementById('main-chart');
        var myChart = echarts.init(chartDom);
        var option;

        // 准备数据
        var dates = {date_list};
        var kData = {k_data};
        var tradingLineData = {trading_line_data};
        var shortLineData = {short_line_data};
        var upCrossIndices = {up_cross_x};
        var upCrossValues = {up_cross_y};
        var downCrossIndices = {down_cross_x};
        var downCrossValues = {down_cross_y};
        var diffData = {diff_data};

        // 配置图表选项
        option = {{
            title: {{
                text: '{symbol} {period_name} 走势与金叉死叉分析',
                left: 'center'
            }},
            tooltip: {{
                trigger: 'axis',
                axisPointer: {{
                    type: 'cross'
                }},
                formatter: function(params) {{
                    var result = params[0].axisValueLabel + '<br/>';
                    params.forEach(function(param) {{
                        var color = param.color;
                        var seriesName = param.seriesName;
                        var value = param.value;
                        var marker = '<span style="display:inline-block;margin-right:5px;border-radius:10px;width:10px;height:10px;background-color:' + color + '"></span>';
                        
                        if (seriesName === 'K线') {{
                            result += marker + seriesName + '<br/>';
                            result += '开盘: ' + value[1] + '<br/>';
                            result += '收盘: ' + value[2] + '<br/>';
                            result += '最低: ' + value[3] + '<br/>';
                            result += '最高: ' + value[4] + '<br/>';
                        }} else {{
                            result += marker + seriesName + ': ' + value + '<br/>';
                        }}
                    }});
                    return result;
                }}
            }},
            legend: {{
                data: ['K线', '操盘线', '空头线', '金叉', '死叉', '差值'],
                top: '30px'
            }},
            grid: [
                {{
                    left: '10%',
                    right: '10%',
                    height: '35%'
                }},
                {{
                    left: '10%',
                    right: '10%',
                    top: '45%',
                    height: '30%'
                }},
                {{
                    left: '10%',
                    right: '10%',
                    top: '80%',
                    height: '15%'
                }}
            ],
            xAxis: [
                {{
                    type: 'category',
                    data: dates,
                    boundaryGap: false,
                    axisLine: {{onZero: false}},
                    splitLine: {{show: false}},
                    min: 'dataMin',
                    max: 'dataMax',
                    axisPointer: {{
                        z: 100
                    }}
                }},
                {{
                    type: 'category',
                    gridIndex: 1,
                    data: dates,
                    boundaryGap: false,
                    axisLine: {{onZero: false}},
                    splitLine: {{show: false}},
                    min: 'dataMin',
                    max: 'dataMax'
                }},
                {{
                    type: 'category',
                    gridIndex: 2,
                    data: dates,
                    boundaryGap: false,
                    axisLine: {{onZero: false}},
                    splitLine: {{show: false}},
                    min: 'dataMin',
                    max: 'dataMax'
                }}
            ],
            yAxis: [
                {{
                    scale: true,
                    splitArea: {{
                        show: true
                    }}
                }},
                {{
                    scale: true,
                    gridIndex: 1,
                    splitNumber: 2,
                    axisLabel: {{show: true}},
                    axisLine: {{show: true}},
                    axisTick: {{show: true}},
                    splitLine: {{show: true}}
                }},
                {{
                    scale: true,
                    gridIndex: 2,
                    splitNumber: 2,
                    axisLabel: {{show: true}},
                    axisLine: {{show: true}},
                    axisTick: {{show: true}},
                    splitLine: {{show: true}}
                }}
            ],
            dataZoom: [
                {{
                    type: 'inside',
                    xAxisIndex: [0, 1, 2],
                    start: 50,
                    end: 100
                }},
                {{
                    show: true,
                    xAxisIndex: [0, 1, 2],
                    type: 'slider',
                    top: '97%',
                    start: 50,
                    end: 100
                }}
            ],
            series: [
                {{
                    name: 'K线',
                    type: 'candlestick',
                    data: kData,
                    itemStyle: {{
                        color: '#ef232a',
                        color0: '#14b143',
                        borderColor: '#ef232a',
                        borderColor0: '#14b143'
                    }}
                }},
                {{
                    name: '操盘线',
                    type: 'line',
                    data: tradingLineData,
                    smooth: true,
                    lineStyle: {{
                        width: 2,
                        color: '#FF0000'
                    }},
                    xAxisIndex: 1,
                    yAxisIndex: 1
                }},
                {{
                    name: '空头线',
                    type: 'line',
                    data: shortLineData,
                    smooth: true,
                    lineStyle: {{
                        width: 2,
                        color: '#0000FF'
                    }},
                    xAxisIndex: 1,
                    yAxisIndex: 1
                }},
                {{
                    name: '差值',
                    type: 'line',
                    data: diffData,
                    areaStyle: {{
                        color: function(params) {{
                            var value = params.data;
                            return value >= 0 ? 'rgba(255, 69, 0, 0.5)' : 'rgba(0, 255, 0, 0.5)';
                        }}
                    }},
                    lineStyle: {{
                        width: 1,
                        color: '#333333'
                    }},
                    xAxisIndex: 2,
                    yAxisIndex: 2
                }},
                {{
                    name: '零线',
                    type: 'line',
                    data: Array(dates.length).fill(0),
                    lineStyle: {{
                        width: 1,
                        type: 'dashed',
                        color: '#000000',
                        opacity: 0.5
                    }},
                    xAxisIndex: 2,
                    yAxisIndex: 2,
                    tooltip: {{
                        show: false
                    }}
                }}
            ]
        }};

        // 添加金叉点
        if (upCrossIndices.length > 0) {{
            var upCrossData = [];
            for (var i = 0; i < upCrossIndices.length; i++) {{
                upCrossData.push([upCrossIndices[i], upCrossValues[i]]);
            }}
            
            option.series.push({{
                name: '金叉',
                type: 'scatter',
                data: upCrossData,
                symbol: 'triangle',
                symbolSize: 15,
                itemStyle: {{
                    color: '#00FF00'
                }},
                xAxisIndex: 1,
                yAxisIndex: 1
            }});
        }}

        // 添加死叉点
        if (downCrossIndices.length > 0) {{
            var downCrossData = [];
            for (var i = 0; i < downCrossIndices.length; i++) {{
                downCrossData.push([downCrossIndices[i], downCrossValues[i]]);
            }}
            
            option.series.push({{
                name: '死叉',
                type: 'scatter',
                data: downCrossData,
                symbol: 'triangle-down',
                symbolSize: 15,
                itemStyle: {{
                    color: '#FF0000'
                }},
                xAxisIndex: 1,
                yAxisIndex: 1
            }});
        }}

        // 使用配置项显示图表
        myChart.setOption(option);
        
        // 响应式调整
        window.addEventListener('resize', function() {{
            myChart.resize();
        }});
    </script>
</body>
</html>
    """
    
    # 保存为HTML文件
    html_filename = f"{symbol.replace('.', '_')}_{period_name}_金叉死叉分析.html"
    try:
        with open(html_filename, 'w', encoding='utf-8') as f:
            f.write(html_content)
        print(f"成功生成HTML文件: {html_filename}")
    except Exception as e:
        print(f"生成HTML文件时出错: {e}")
        import traceback
        traceback.print_exc()
    
    return html_filename

# 主函数
print("开始验证均线上穿下穿算法...")

# 周期名称，用于图表标题
period_name = "日线"  # 周期名称，用于图表标题

# 允许用户自定义数据的时间范围
def get_user_params():
    """获取用户自定义的参数，包括日期范围和K线数量"""
    parser = argparse.ArgumentParser(description='均线金叉死叉算法验证工具')
    parser.add_argument('--start', type=str, help='起始日期，格式：YYYY-MM-DD')
    parser.add_argument('--end', type=str, help='结束日期，格式：YYYY-MM-DD')
    parser.add_argument('--years', type=float, default=DEFAULT_YEARS, 
                        help=f'获取最近几年的数据，默认为{DEFAULT_YEARS}年')
    parser.add_argument('--max-klines', type=int, default=MAX_KLINES_TO_DISPLAY, 
                        help=f'最大显示K线数量，默认为{MAX_KLINES_TO_DISPLAY}根')
    parser.add_argument('--symbol', type=str, default="sh000001", 
                        help='股票或指数代码，默认为上证指数(sh000001)')
    args = parser.parse_args()
    
    # 默认结束日期为当前日期
    today = datetime.now().date()
    end_date = today
    
    # 如果指定了结束日期，则使用指定的结束日期
    if args.end:
        try:
            end_date = pd.Timestamp(args.end).date()
        except:
            print(f"警告：无效的结束日期格式 '{args.end}'，使用默认值 {end_date}")
    
    # 使用用户指定的年数
    years = args.years
    
    # 计算默认的起始日期
    default_start_date = end_date - relativedelta(years=years)
    
    # 如果指定了起始日期，则使用指定的起始日期
    start_date = default_start_date
    if args.start:
        try:
            start_date = pd.Timestamp(args.start).date()
        except:
            print(f"警告：无效的起始日期格式 '{args.start}'，使用默认值 {start_date}")
    
    # 获取最大K线数量
    max_klines = args.max_klines
    print(f"设置最大显示K线数量: {max_klines}根")
    print(f"设置数据年限: {years}年")
    
    # 获取股票代码
    symbol = args.symbol
    print(f"使用股票/指数代码: {symbol}")
    
    return start_date, end_date, max_klines, symbol

try:
    # 获取用户自定义的参数
    start_date, cutoff_date, max_klines, symbol = get_user_params()
    
    # 使用 akshare 获取 K 线数据
    # 判断是股票还是指数
    if symbol.startswith("sh") or symbol.startswith("sz"):
        # 指数数据
        if symbol.startswith("sh") and len(symbol) == 8:
            print(f"获取上海指数 {symbol} 的日线数据")
            klines_df = ak.stock_zh_index_daily(symbol=symbol)
        elif symbol.startswith("sz") and len(symbol) == 8:
            print(f"获取深圳指数 {symbol} 的日线数据")
            klines_df = ak.stock_zh_index_daily(symbol=symbol)
        else:
            # 股票数据
            print(f"获取股票 {symbol} 的日线数据")
            klines_df = ak.stock_zh_a_daily(symbol=symbol)
    
    print(f"使用截止日期: {cutoff_date}")
    print(f"使用起始日期: {start_date}")
    
    # 筛选指定日期范围的数据
    klines_df = klines_df[(klines_df['date'] <= cutoff_date) & (klines_df['date'] >= start_date)]
    
    # 如果数据量超过最大显示数量，则只保留最近的数据
    if len(klines_df) > max_klines:
        print(f"数据量({len(klines_df)})超过最大显示数量({max_klines})，将只显示最近的{max_klines}根K线")
        klines_df = klines_df.sort_values('date', ascending=False).head(max_klines).sort_values('date')
    
    print(f"筛选后的数据量: {len(klines_df)} 条记录")
    
    # 重命名列以匹配代码中使用的列名
    klines_df = klines_df.rename(columns={
        'date': 'datetime',  # 日期列重命名为datetime
        'open': 'open',      # 开盘价
        'high': 'high',      # 最高价
        'low': 'low',        # 最低价
        'close': 'close'     # 收盘价
    })
    
    # 转换为类似于原代码中使用的对象格式
    class KLineData:
        def __init__(self, df):
            # 直接使用DataFrame的Series，这样可以使用iloc属性
            self.datetime = df['datetime']
            self.open = df['open']
            self.high = df['high']
            self.low = df['low']
            self.close = df['close']
            self._length = len(df)
            self._df = df  # 保存原始DataFrame以便后续使用
        
        def __len__(self):
            return self._length
    
    klines = KLineData(klines_df)
    
    # 检查数据是否为空
    if len(klines) == 0:
        print("错误：筛选后的数据为空，请调整日期范围")
        exit(1)
    
    # 调试输出，查看日期格式
    sample_ts = klines.datetime.iloc[0]  # 使用iloc而不是索引访问
    print(f"调试信息 - 日期样本: {sample_ts}, 类型: {type(sample_ts)}")
    
    # 计算信号
    result = calculate_trading_signals(klines)
    
    # 生成HTML文件
    html_filename = generate_html(result, symbol, period_name)
    
    # 输出结果
    if html_filename:
        print(f"文件路径: {html_filename}")
        print(f"请在浏览器中打开此文件查看K线图和双均线金叉死叉分析图")
        print("-"*80)
        print(f"【信号统计】")
        print(f"金叉信号数量: {result['signal'].value_counts().get('上穿', 0)}")
        print(f"死叉信号数量: {result['signal'].value_counts().get('下穿', 0)}")
        print(f"总信号数量: {result['signal'].notna().sum()}")
        print(f"信号占比: {(result['signal'].notna().sum() / len(result) * 100):.2f}%")
        print("="*80)
    
    print("已使用原生ECharts生成增强版可视化图表，包含K线图和双均线金叉死叉标注")
    
except Exception as e:
    print(f"发生错误: {e}")
    # 打印更详细的错误信息
    import traceback
    traceback.print_exc()
finally:
    print("处理完成")